import re
import os
import csv
from datetime import datetime

def extract_five_tuple_and_cid_from_txt(file_path):
    """
    从txt文件中提取五元组、cid和完整的HTTP请求信息
    """
    results = []

    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
            
        # 查找Header部分的JSON数据
        header_pattern = r'Header:\{([^}]+)\}'
        header_matches = re.findall(header_pattern, content)
        
        # 查找包含cyber关键字的内容
        cyber_matches = re.findall(r'cyber', content, re.IGNORECASE)
        
        # 查找Access-Control-Request-Method值
        method_pattern = r'Access-Control-Request-Method:\s*([A-Z]+)'
        method_matches = re.findall(method_pattern, content)
        
        # 查找完整的HTTP请求行 (例如: GET /api/... HTTP/1.1)
        http_request_pattern = r'(GET|POST|PUT|DELETE|HEAD|OPTIONS|PATCH)\s+([^\s]+)\s+HTTP/[\d.]+'
        http_requests = re.findall(http_request_pattern, content)
        
        # 查找HTTP方法 (GET, POST, OPTIONS等) - 更宽泛的匹配
        http_method_pattern = r'(GET|POST|PUT|DELETE|HEAD|OPTIONS|PATCH)\s+/.*cyber'
        http_methods = re.findall(http_method_pattern, content)
        
        for i, header in enumerate(header_matches):
            # 从Header中提取各个字段
            cid_match = re.search(r'"cid":"([^"]*)"', header)
            sip_match = re.search(r'"sip":"([^"]*)"', header)
            spt_match = re.search(r'"spt":(\d+)', header)
            dip_match = re.search(r'"dip":"([^"]*)"', header)
            dpt_match = re.search(r'"dpt":(\d+)', header)
            proto_match = re.search(r'"proto":"([^"]*)"', header)
            
            # 获取对应的Access-Control-Request-Method值（如果存在）
            access_method = method_matches[i] if i < len(method_matches) else ''
            
            # 获取对应的HTTP方法（如果存在）
            http_method = http_methods[i] if i < len(http_methods) else ''
            
            # 获取完整的HTTP请求行（如果存在）
            full_http_request = ''
            if i < len(http_requests):
                full_http_request = f"{http_requests[i][0]} {http_requests[i][1]} HTTP/1.1"
            
            result = {
                'cid': cid_match.group(1) if cid_match else '',
                'sip': sip_match.group(1) if sip_match else '',
                'spt': spt_match.group(1) if spt_match else '',
                'dip': dip_match.group(1) if dip_match else '',
                'dpt': dpt_match.group(1) if dpt_match else '',
                'proto': proto_match.group(1) if proto_match else '',
                'type': access_method,  # Access-Control-Request-Method
                'http_method': http_method,  # HTTP方法
                'full_http_request': full_http_request,  # 完整的HTTP请求行
                'cyber_count': len(cyber_matches),  # cyber关键字出现次数
                'filename': os.path.basename(file_path)
            }
            results.append(result)

    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {e}")

    return results

def process_all_txt_files(folder_path):
    """
    遍历文件夹中的所有txt文件并提取信息
    """
    all_results = []

    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"文件夹 {folder_path} 不存在")
        return all_results

    # 遍历文件夹中的所有文件
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.txt'):
                file_path = os.path.join(root, file)
                print(f"正在处理文件: {file_path}")
                results = extract_five_tuple_and_cid_from_txt(file_path)
                if results:  # 只有当有结果时才添加
                    all_results.extend(results)

    return all_results

def save_to_csv(data, output_file):
    """
    将提取的数据保存到CSV文件
    """
    if not data:
        print("没有数据需要保存")
        return

    with open(output_file, 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['filename', 'cid', 'sip', 'spt', 'dip', 'dpt', 'proto', 'type', 'http_method', 'full_http_request', 'cyber_count']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        writer.writeheader()
        for item in data:
            writer.writerow(item)

    print(f"数据已保存到 {output_file}")

def main():
    # 源文件夹路径
    folder_path = r"D:\DPI测试\转储测试\Data"
    # 确保输出目录存在
    output_dir = r"D:\DPI测试\转储测试\Data"
    os.makedirs(output_dir, exist_ok=True)
    output_csv_path = os.path.join(output_dir, f"kafka_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv")

    # 处理所有txt文件
    print(f"正在遍历文件夹: {folder_path}")
    all_results = process_all_txt_files(folder_path)

    if all_results:
        print(f"总共找到 {len(all_results)} 条记录:")

        # 显示前几条记录作为示例
        for i, record in enumerate(all_results[:5]):
            print(f"{i + 1}. 文件:{record['filename']} CID:{record['cid']} "
                  f"SIP:{record['sip']} SPT:{record['spt']} DIP:{record['dip']} "
                  f"DPT:{record['dpt']} PROTO:{record['proto']} TYPE:{record['type']} "
                  f"HTTP_METHOD:{record['http_method']} FULL_HTTP_REQUEST:{record['full_http_request']} "
                  f"CYBER_COUNT:{record['cyber_count']}")

        if len(all_results) > 5:
            print("...")

        # 保存到CSV文件
        save_to_csv(all_results, output_csv_path)
        print(f"提取完成，共处理 {len(all_results)} 条记录")
    else:
        print("未找到任何五元组信息")

if __name__ == "__main__":
    main()